#include <iostream>
#include <chrono>
#include <onnxruntime_cxx_api.h>
#include <cpu_provider_factory.h>

// namespace std;
// namespace Ort;

int main() {
    // 初始化环境
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "hcn");

    // 创建会话配置
    Ort::SessionOptions session_options;
    Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CPU(session_options, 1));
    Ort::Session session(env, "../hcn2.onnx", session_options);

    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
    Ort::AllocatorWithDefaultOptions allocator;


    // 创建输入tensor
    float input_data1 = 9;
    float input_data2[4][1][1][96] = {0};
    std::vector<int64_t> input_shape1 = {1,};
    std::vector<int64_t> input_shape2 = {4, 1, 1, 96};

    std::vector<Ort::Value> input_tensors;
    // Ort::Value input_tensor1 = Ort::Value::CreateTensor<float>(memory_info, &input_data1, 1, input_shape1.data(), input_shape1.size());
    // Ort::Value input_tensor2 = Ort::Value::CreateTensor<float>(memory_info, &input_data2[0][0][0][0], 4*96, input_shape2.data(), input_shape2.size());
    // input_tensors.push_back(std::move(input_tensor1));
    // input_tensors.push_back(std::move(input_tensor2));
    input_tensors.push_back(Ort::Value::CreateTensor<float>(memory_info, &input_data1, 1, input_shape1.data(), input_shape1.size()));
    input_tensors.push_back(Ort::Value::CreateTensor<float>(memory_info, &input_data2[0][0][0][0], 4*96, input_shape2.data(), input_shape2.size()));

    // std::vector<const char*> input_names = {"x","hc"};
    std::vector<const char*> input_names;
    auto input_name1 = session.GetInputNameAllocated(0,allocator);
    auto input_name2 = session.GetInputNameAllocated(1,allocator);
    input_names.push_back(input_name1.get());
    input_names.push_back(input_name2.get());
    std::cout << "input_name1: " << input_name1.get() << std::endl;
    std::cout << "input_name2: " << input_name2.get() << std::endl;

    // std::cout << "=========================hc_data=====================================" << std::endl;
    // void* hc_data = input_tensors[1].GetTensorMutableData<void>();
    // float* hc_floatData = static_cast<float*>(hc_data);
    // for (size_t i = 0; i < input_tensors[1].GetTensorTypeAndShapeInfo().GetElementCount(); i++) {
    //   std::cout << hc_floatData[i] << " ";
    // }
    // std::cout << "==============================================================" << std::endl;

    // output
    float output_data1[1][250];
    float output_data2[4][1][1][96];
    std::vector<int64_t> output_shape1 = {1,250};
    std::vector<int64_t> output_shape2 = {4, 1, 1, 96};

    std::vector<Ort::Value> output_tensors;
    // Ort::Value output_tensor1 = Ort::Value::CreateTensor<float>(memory_info, &output_data1[0][0], 250, output_shape1.data(), output_shape1.size());
    // Ort::Value output_tensor2 = Ort::Value::CreateTensor<float>(memory_info, &output_data2[0][0][0][0], 4*96, output_shape2.data(), output_shape2.size());
    // output_tensors.push_back(std::move(output_tensor1));
    // output_tensors.push_back(std::move(output_tensor2));
    output_tensors.push_back(Ort::Value::CreateTensor<float>(memory_info, &output_data1[0][0], 250, output_shape1.data(), output_shape1.size()));
    output_tensors.push_back(Ort::Value::CreateTensor<float>(memory_info, &output_data2[0][0][0][0], 4*96, output_shape2.data(), output_shape2.size()));

    // std::vector<const char*> output_names = {"164","167"};
    std::vector<const char*> output_names;
    auto output_name1 = session.GetOutputNameAllocated(0,allocator);
    auto output_name2 = session.GetOutputNameAllocated(1,allocator);
    output_names.push_back(output_name1.get());
    output_names.push_back(output_name2.get());
    std::cout << "output_name1: " << output_name1.get() << std::endl;
    std::cout << "output_name2: " << output_name2.get() << std::endl;


    // 运行模型
    // Ort::RunOptions run_options;
    Ort::RunOptions run_options = Ort::RunOptions{ nullptr };

    session.Run(run_options, input_names.data(), input_tensors.data(), 2, output_names.data(), output_tensors.data(), 2);
    // std::vector<Ort::Value> output_tensors = session.Run(run_options, input_names.data(), input_tensors.data(), 2, output_names.data(), 2);
    // output_tensors = session.Run(run_options, input_names.data(), input_tensors.data(), 2, output_names.data(), 2);

    std::cout << "output_tensors.size(): " << output_tensors.size() << std::endl;
    std::cout << "std::move(output_tensors[0]): " << std::move(output_tensors[0]) << std::endl;
    std::cout << "std::move(output_tensors[1]): " << std::move(output_tensors[1]) << std::endl;

    std::cout << "===========================output_tensors[0]===================================" << std::endl;
    std::cout << "output_tensors[0].IsTensor(): " << output_tensors[0].IsTensor() << std::endl;
    std::cout << "output_tensors[0].GetTensorTypeAndShapeInfo(): " << output_tensors[0].GetTensorTypeAndShapeInfo() << std::endl;
    std::vector<int64_t> x = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape();
    std::cout << "output_tensors[0].GetTensorTypeAndShapeInfo().GetShape(): ";
    for (int64_t num:x) {
      std::cout << num << " ";
    }
    std::cout << std::endl;
    std::cout << "output_tensors[0].GetTensorTypeAndShapeInfo().GetDimensionsCount(): " << output_tensors[0].GetTensorTypeAndShapeInfo().GetDimensionsCount() << std::endl;
    std::cout << "output_tensors[0].GetTensorTypeAndShapeInfo().GetElementType(): " << output_tensors[0].GetTensorTypeAndShapeInfo().GetElementType() << std::endl;
    std::cout << "output_tensors[0].GetTensorTypeAndShapeInfo().GetElementCount(): " << output_tensors[0].GetTensorTypeAndShapeInfo().GetElementCount() << std::endl;
    std::cout << "output_tensors[0].GetTensorMutableData<float>(): " << output_tensors[0].GetTensorMutableData<float>() << std::endl;


    std::cout << "===========================output_tensors[1]===================================" << std::endl;
    std::cout << "output_tensors[1].IsTensor(): " << output_tensors[1].IsTensor() << std::endl;
    std::cout << "output_tensors[1].GetTensorTypeAndShapeInfo(): " << output_tensors[1].GetTensorTypeAndShapeInfo() << std::endl;
    x = output_tensors[1].GetTensorTypeAndShapeInfo().GetShape();
    std::cout << "output_tensors[1].GetTensorTypeAndShapeInfo().GetShape(): ";
    for (int64_t num:x) {
      std::cout << num << " ";
    }
    std::cout << std::endl;
    std::cout << "output_tensors[1].GetTensorTypeAndShapeInfo().GetDimensionsCount(): " << output_tensors[1].GetTensorTypeAndShapeInfo().GetDimensionsCount() << std::endl;
    std::cout << "output_tensors[1].GetTensorTypeAndShapeInfo().GetElementType(): " << output_tensors[1].GetTensorTypeAndShapeInfo().GetElementType() << std::endl;
    std::cout << "output_tensors[1].GetTensorTypeAndShapeInfo().GetElementCount(): " << output_tensors[1].GetTensorTypeAndShapeInfo().GetElementCount() << std::endl;
    std::cout << "output_tensors[1].GetTensorMutableData<float>(): " << output_tensors[1].GetTensorMutableData<float>() << std::endl;

    std::cout << "===========================output_tensors[0]===================================" << std::endl;
    void* data = output_tensors[0].GetTensorMutableData<void>();
    float* floatData = static_cast<float*>(data);
    for (size_t i = 0; i < output_tensors[0].GetTensorTypeAndShapeInfo().GetElementCount(); i++) {
      std::cout << floatData[i] << " ";
      // printf("%0.8f ", floatData[i]);
    }
    std::cout << "==============================================================" << std::endl;

    std::cout << "===========================output_tensors[1]===================================" << std::endl;
    float hc[4][1][1][96]={0};
    size_t num = 96;
    data = output_tensors[1].GetTensorMutableData<void>();
    floatData = static_cast<float*>(data);
    for (size_t i = 0; i < output_tensors[1].GetTensorTypeAndShapeInfo().GetElementCount(); i++) {
      std::cout << floatData[i] << ", ";
      if ((i+1)%num==0) std::cout << " ][ ";
      hc[i/num][0][0][i%num] = std::round(floatData[i] * 1000000) / 1000000;
    }
    std::cout << "==============================================================" << std::endl;

    std::cout << "============================hc==================================" << std::endl;
    std::vector<int64_t> input_shape22 = {4, 1, 1, 96};
    Ort::Value input_tensor3 = Ort::Value::CreateTensor<float>(memory_info, &hc[0][0][0][0], 4*96, input_shape22.data(), input_shape22.size());
    std::cout << hc << std::endl;
    float* hc_floatData = input_tensor3.GetTensorMutableData<float>();
    for (size_t i = 0; i < input_tensor3.GetTensorTypeAndShapeInfo().GetElementCount(); i++) {
      std::cout << hc_floatData[i] << " ";
    }
    std::cout << "==============================================================" << std::endl;

    
    float tf = 10;
    std::vector<Ort::Value> input_tensors2;
    std::vector<int64_t> input_shape12 = {1,};
    // Ort::Value input_tensorx = Ort::Value::CreateTensor<float>(memory_info, &tf, 1, input_shape12.data(), 1);
    // input_tensors2.push_back(std::move(input_tensorx));
    input_tensors2.push_back(Ort::Value::CreateTensor<float>(memory_info, &tf, 1, input_shape12.data(), 1));
    input_tensors2.push_back(std::move(output_tensors[1]));
    // input_tensors2.push_back(std::move(input_tensor3));

    float output_data12[1][250];
    float output_data22[4][1][1][96];
    std::vector<int64_t> output_shape12 = {1,250};
    std::vector<int64_t> output_shape22 = {4, 1, 1, 96};

    std::vector<Ort::Value> output_tensors2;
    // Ort::Value output_tensor12 = Ort::Value::CreateTensor<float>(memory_info, &output_data12[0][0], 250, output_shape12.data(), output_shape12.size());
    // Ort::Value output_tensor22 = Ort::Value::CreateTensor<float>(memory_info, &output_data22[0][0][0][0], 4*96, output_shape22.data(), output_shape22.size());
    // output_tensors2.push_back(std::move(output_tensor12));
    // output_tensors2.push_back(std::move(output_tensor22));
    output_tensors2.push_back(Ort::Value::CreateTensor<float>(memory_info, &output_data12[0][0], 250, output_shape12.data(), output_shape12.size()));
    output_tensors2.push_back(Ort::Value::CreateTensor<float>(memory_info, &output_data22[0][0][0][0], 4*96, output_shape22.data(), output_shape22.size()));

    std::vector<const char*> input_names2;
    auto input_name12 = session.GetInputNameAllocated(0,allocator);
    auto input_name22 = session.GetInputNameAllocated(1,allocator);
    input_names2.push_back(input_name12.get());
    input_names2.push_back(input_name22.get());
    std::cout << "input_name12: " << input_name12.get() << std::endl;
    std::cout << "input_name22: " << input_name22.get() << std::endl;

    std::vector<const char*> output_names2;
    auto output_name12 = session.GetOutputNameAllocated(0,allocator);
    auto output_name22 = session.GetOutputNameAllocated(1,allocator);
    output_names2.push_back(output_name12.get());
    output_names2.push_back(output_name22.get());
    std::cout << "output_name12: " << output_name12.get() << std::endl;
    std::cout << "output_name22: " << output_name22.get() << std::endl;

    // std::vector<Ort::Value> output_tensors2 = session.Run(run_options, input_names.data(), input_tensors2.data(), 2, output_names.data(), 2);
    // std::vector<Ort::Value> output_tensors2;
    session.Run(run_options, input_names2.data(), input_tensors2.data(), 2, output_names2.data(), output_tensors2.data(), 2);
    std::cout << "output_tensors.size(): " << output_tensors2.size() << std::endl;

    std::cout << "============================output_tensors2[0]==================================" << std::endl;
    // void* data2 = output_tensors2[0].GetTensorMutableData<void>();
    // float* floatData2 = static_cast<float*>(data2);
    float* floatData2 = output_tensors2[0].GetTensorMutableData<float>();
    for (size_t i = 0; i < output_tensors2[0].GetTensorTypeAndShapeInfo().GetElementCount(); i++) {
      std::cout << floatData2[i] << " ";
      // printf("%0.8f ", floatData[i]);
    }
    std::cout << "==============================================================" << std::endl;

    std::cout << "==========================output_tensors2[1]====================================" << std::endl;
    data = output_tensors2[1].GetTensorMutableData<void>();
    floatData = static_cast<float*>(data);
    for (size_t i = 0; i < output_tensors2[1].GetTensorTypeAndShapeInfo().GetElementCount(); i++) {
      std::cout << floatData[i] << " ";
    }
    std::cout << "==============================================================" << std::endl;
    

    float itValue = 0.1;
    auto start = std::chrono::high_resolution_clock::now();
    for (int i=0; i<10000; ++i) {
        itValue += 0.001;

        // input_tensor1 = Ort::Value::CreateTensor<float>(memory_info, &itValue, 1, input_shape1.data(), 1);
        input_tensors.push_back(Ort::Value::CreateTensor<float>(memory_info, &itValue, 1, input_shape1.data(), 1));
        input_tensors.push_back(std::move(output_tensors[1]));
        output_tensors = session.Run(run_options, input_names.data(), input_tensors.data(), 2, output_names.data(), 2);
    }
    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed = end - start;
    std::cout << "Elapsed time: " << elapsed.count() << " seconds" << std::endl;




    return 0;
}
